#!/usr/bin/env python
# -*- coding:utf-8 -*-
from __future__ import absolute_import

__author__ = 'wshu'
__version__ = '1.0'
"""
    ***********************************
    *  @filename : tasks.py
    *  @Author : wshu
    *  @CodeDate : 2020/3/17 10:46
    *  @Software : PyCharm
    ***********************************
    扫描器集成管理(试用版未开放此功能)
    支持(awvs/nessus/nikto)
"""
import time
from celery import shared_task
from celery.utils.log import get_task_logger
from django.core.cache import cache

from TaskManage.core import awvs, nessus
from contextlib import contextmanager

from JQSetting.core import Awvs, Nessus

from JQSetting.comm.mails import send_notice_mail

from TaskManage.models import Task
from NoticeManage.views import notice_add

logger = get_task_logger(__name__)

LOCK_EXPIRE = 60 * 10  # Lock expires in 10 minutes


@shared_task
def awvs_store(scanid, taskid):
    """
    :param scanid:  扫描id
    :param taskid:  任务id
    :return:
    """
    task = Task.objects.filter(task_id =taskid ).first()
    while True:
        res = Nessus.details(scanid,task.task_scanner.id)
        try:
            res['info']['status']
        except:
            continue
        if res['info']['status'] == 'canceled' or res['info']['status'] == 'completed':
            #time.sleep(600)
            nessus.get_scan_vuln(scanid,task,task.task_scanner.id)
            task.task_status=4
            task.save()
            data={
                  'notice_title':'任务进度通知',
                  'notice_body':'您对'+task.task_name+'的扫描任务已完成，请及时查看结果',
                  'notice_url':'/task/user/',
                  'notice_type':'notice',
                  }
            user = task.task_user
            notice_add(user,data)
            send_notice_mail(user.email,data)
            break
        else:
            time.sleep(30)



@shared_task
def nessus_store(scanid, taskid):
    """
    :param scanid: 扫描id
    :param taskid: 任务id
    :return:
    """
    task = Task.objects.filter(task_id =taskid ).first()
    while True:
        status = Awvs.getstatus(scanid,task.task_scanner.id)
        if status == 'completed':
            awvs.get_scan_result(scanid,taskid,task.task_scanner.id)
            task.task_status=4
            task.save()
            #type_task_list = {'移动应用':'type1','web应用':'type2','操作系统':'type3'}
            data={
                  'notice_title':'任务进度通知',
                  'notice_body':'您对'+task.task_name+'的扫描任务已完成，请及时查看结果',
                  'notice_url':'/task/user/',
                  'notice_type':'notice',
                  }
            user = task.task_user
            notice_add(user,data)
            send_notice_mail(user.email,data)
            break
        elif status == 'aborted':
            awvs.get_scan_result(scanid,taskid,task.task_scanner.id)
            task.task_status=5
            task.save()
            data={
                  'notice_title':'任务进度通知',
                  'notice_body':'您对'+task.task_name+'的扫描任务已完成，请及时查看结果',
                  'notice_url':'/task/user/',
                  'notice_type':'notice',
                  }
            user = task.task_user
            notice_add(user,data)
            send_notice_mail(user.email,data)
            break
        else:
            time.sleep(60)


@shared_task
def nikto_store(scanid, taskid):
    """
    :param scanid: 扫描id
    :param taskid: 任务id
    :return:
    """